package com.vipboot.framework.core.util.mdc;

import jakarta.validation.constraints.NotNull;
import org.slf4j.MDC;
import org.springframework.core.task.VirtualThreadTaskExecutor;

import java.util.concurrent.Callable;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Future;

public class MdcVirtualThreadTaskExecutor extends VirtualThreadTaskExecutor {
    public MdcVirtualThreadTaskExecutor(String threadNamePrefix) {
        super(threadNamePrefix);
    }

    @Override
    public void execute(@NotNull Runnable task) {
        super.execute(MdcThreadWrapUtil.wrap(task, MDC.getCopyOfContextMap()));
    }

    @Override
    public <T> Future<T> submit(@NotNull Callable<T> task) {
        return super.submit(MdcThreadWrapUtil.wrap(task, MDC.getCopyOfContextMap()));
    }

    @Override
    public Future<?> submit(@NotNull Runnable task) {
        return super.submit(MdcThreadWrapUtil.wrap(task, MDC.getCopyOfContextMap()));
    }

    @Override
    public CompletableFuture<Void> submitCompletable(Runnable task) {
        return super.submitCompletable(MdcThreadWrapUtil.wrap(task, MDC.getCopyOfContextMap()));
    }

    @Override
    public <T> CompletableFuture<T> submitCompletable(Callable<T> task) {
        return super.submitCompletable(MdcThreadWrapUtil.wrap(task, MDC.getCopyOfContextMap()));
    }
}
